[KMCompiler][ttx] Optimize silu with rowwise nomask kernels#365
[KMCompiler][ttx] Optimize silu with rowwise nomask kernels#365YangLong114514 wants to merge 2 commits into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request optimizes the SiLU activation implementation for NPU by introducing autotuned, non-masked forward kernels (_silu_fwd_nomask_kernel and _silu_fwd_nomask_single_kernel) along with helper functions for dynamic grid and block size calculation. The review feedback points out a critical issue where an empty tensor (with zero rows) would result in an invalid grid size of 0, causing a runtime error. Additionally, it recommends reusing the cached get_num_cores utility from .utils instead of querying NPU device properties directly on every kernel launch to eliminate unnecessary CPU overhead.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
YangLong114514
changed the title
1 participant
Description
The silu operator has been optimized for the Ascend platform.
Changes
Updated the
BLOCK_SIZE_Ncalculation to use a capped power-of-two tile size, reducing unstable tile-size jumps around boundary shapes.Added dynamic grid calculation based on the autotuned
BLOCK_SIZE_M, avoiding unnecessary empty programs for small shapes.Added no-mask forward kernels for shapes that are fully divisible by the row and column tiles, removing mask construction and masked
tl.load/tl.storeoverhead.Added a no-mask single-tile forward kernel for shapes where
n_cols == BLOCK_SIZE_N, eliminating the inner column loop for common benchmark shapes.Performance
Using Ascend 910C and Triton 3.2.x of FlagTree:
Accuracy test
The correctness verification of the SiLU operator for all three data types, float32, float16, and bf16, passed in the Mojo correctness verification directory.
